import gymnasium as gym
import ale_py
import keyboard
import time

gym.register_envs(ale_py)

# 定义键盘按键与动作的映射
key_action_map = {
    'left': 3,  # 左移
    'right': 2,  # 右移
    'space': 1  # 发射小球
}

def handle_keyboard_input():
    for key, action in key_action_map.items():
        if keyboard.is_pressed(key):
            return action
    return 0  # 默认动作：不做任何操作

env = gym.make('ALE/Breakout-v5', render_mode="human")
observation, info = env.reset()
done = False

while not done:
    # 获取键盘输入对应的动作
    action = handle_keyboard_input()
    # 在环境中执行动作
    observation, reward, terminated, truncated, info = env.step(action)
    done = terminated or truncated
    time.sleep(0.01)  # 控制游戏速度

env.close()